#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   lxc.py
@Time    :   2024-06-30 14:07
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   对PVE CT容器进行定时开关、重启操作
"""
import logging
import os
import subprocess
import sys
from argparse import ArgumentParser
from os import path
from sys import exit
import tomllib
from datetime import datetime, timedelta
from time import sleep

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


def list_files(directory):
	"""
	列出指定目录下的所有文件
	"""
	file_list = []
	# 使用glob模块获取指定目录下的所有文件
	for root, dirs, files in os.walk(directory):
		for name in files:
			# print(os.path.join(root, name))
			file_list.append(os.path.join(root, name))
	return file_list


class Lxc:
	def __init__(self, config_restart_dir, config_list: str, debug: bool = False):
		"""
		配置文件
		:param config_restart_dir:
		"""
		self.toml_end_time = None
		self.toml_start_time = None
		self.toml_action = None
		self.toml_interval = 0
		self.toml_lxc_list = None
		self.config_list = config_list
		self.default_config = config_restart_dir
		self.debug = debug
		self.config_restart_dir = config_restart_dir
		self.txt_config = {}
		self.txt_config_list = {}
		self.log_file = None
		# 当前时间:分
		self.current_min = 0
		# 当前时间:时
		self.current_hour = 0
		if self.debug:
			logger.debug("开始实例化")
		# 设置实例配置
		self.toml_name = ""
		self.toml_id = 220
		# 开启操作
		self.toml_on_hour = 0
		self.toml_on_delayed = 0
		self.toml_enabled = False
		# 关闭操作
		self.toml_off_hour = 0
		self.toml_off_delayed = 0

	def get_time(self):
		"""
		获取当前时间
		:return:
		"""
		if self.debug:
			logger.debug("正在获取当前时间")
		current = datetime.now()
		self.current_hour = current.hour
		self.current_min = current.minute
		if self.debug:
			logger.debug(f"当前小时: {self.current_hour} - 当前分钟: {self.current_min}")

	def read(self):
		"""
		读取配置文件
		:return:
		"""
		if self.debug:
			logger.debug("正在读取配置文件")
		try:
			with open(self.config_restart_dir, "rb") as f:
				self.txt_config = tomllib.load(f)
				return
		except Exception as e:
			print(self.config_restart_dir)
			logger.error(f"读取配置文件失败: {e}")
			exit(1)

	def read_list(self, file):
		"""
		读取配置文件
		:param file: 需要读取的主机配置文件
		:return:
		"""
		if self.debug:
			logger.debug(f"正在读取配置文件: {file}")
		try:
			with open(file, "rb") as f:
				self.txt_config_list = tomllib.load(f)
				return
		except Exception as e:
			print(file)
			logger.error(f"读取配置文件失败: {e}")
			exit(1)

	def result_log(self, result: str):
		"""
		输出结果
		:return:
		"""
		if self.debug:
			logger.debug("正在输出结果")
		if self.log_file:
			logger.info(f"日志文件: {self.log_file}")
		else:
			logger.info("日志文件: None")
			return
		if not path.isfile(self.log_file):
			os.system(f"touch {self.log_file}")
		try:
			with open(file=self.log_file, mode='a+', encoding='utf-8') as w:
				w.write(result)
		except Exception as e:
			logger.error(f"写入日志文件失败: {e}")

	def ct_op(self, Id: int, op: str):
		"""
		pve开机
		:param op: 操作类型
		:param Id: ID
		:return:
		"""
		cmd_op = "start"
		ops = "开机"
		if op.lower() == "off":
			cmd_op = "stop"
			ops = "关机"

		if self.debug:
			logger.debug(f"正在{ops}: {Id}")
		if os.system(f"lxc-{cmd_op} {Id}") == 0:
			logger.info(f"{ops}成功: {Id}")
			self.result_log(f"{datetime.now()} - {Id} - {ops}成功\n")
		else:
			logger.error(f"{ops}失败: {Id}")
			self.result_log(f"{datetime.now()} - {Id} - {ops}失败\n")

	def get_config_final_time(self, hour, delayed: int):
		"""
		获取配置文件最终时间并计算是否在配置时间内
		:param hour: 配置文件设置的时间
		:param delayed: 配置文件设置的超时时间
		:return: bool(是否需要执行任务)
		"""
		today = datetime.now().date()
		conf_time = datetime.combine(today, datetime.min.time().replace(hour=hour, minute=0, second=0, microsecond=0))
		max_time = conf_time + timedelta(hours=delayed)
		now = datetime.now()
		# 如果当前时间大于配置文件+超时时间,则表示该配置已超时
		if self.debug:
			logger.debug(f"当前时间: {now} - 配置文件时间: {conf_time} - 最大时间: {max_time}")
		if conf_time < now < max_time:
			return True
		else:
			return False

	def operation(self, sec: str, info: dict):
		"""
		执行操作
		:return:
		"""
		try:
			# 获取base信息
			if str(sec).lower() == "list":
				self.toml_lxc_list = info['list']
				self.toml_interval = info['interval']
				self.toml_action = info['action']
				self.toml_enabled = info['enabled']
			else:
				if str(sec).lower() == "restart".lower():
					if str(sec).lower() == "on".lower():
						self.toml_lxc_list = info['list']
						self.toml_interval = info['interval']
						self.toml_start_time = info['start_time']
						self.toml_end_time = info['end_time']
						self.toml_enabled = info['enabled']
					else:
						logger.error(f"未知配置: {sec}")
						exit(1)

				else:
					logger.error(f"未知配置: {sec}")
					exit(1)
		except Exception as e:
			logger.error(f"获取配置参数失败: {e}")
			print(info)
			exit(1)

	def restart(self):
		"""
		重启配置文件
		:return:
		"""
		cl = [self.default_config]
		if path.isdir(args_conf_restart):
			if self.debug:
				logger.debug(f'正在获取配置文件: {args_conf_restart}')
			for i in list_files(args_conf_restart):
				if self.debug:
					logger.debug(f"正在判断: {i}")
				ff = str(i).split(".")[-1]
				if ff == "toml":
					cl.append(path.join(args_conf_restart, i))
		if self.debug:
			logger.debug(f"当前配置文件数量{len(cl)}:")
		for self.config_restart_dir in cl:
			self.toml_enabled = False
			if self.debug:
				logger.debug(f"正在读取配置文件: {self.config_restart_dir}")
			self.read()
			self.get_time()
			if len(self.txt_config) == 0:
				logger.error("配置文件为空")
				exit(1)
			for i in self.txt_config:
				if self.debug:
					logger.debug(self.txt_config[i])
				if str(i) != "restart":
					continue

				self.operation(info=self.txt_config[i], sec=i)

		sleep(5)

	def pve_op_ct_list(self, Id, hour, delayed: int, action: str):
		"""
		执行列表主机操作
		:param Id: CT ID
		:param hour: 配置文件设置的时间
		:param delayed: 配置文件设置的超时时间
		:param action: 配置文件设置的操作
		:return:
		"""
		# 设置两种操作对应的值
		op_dict = {"off": "stop", "on": "start"}
		op_dict_name = {"off": "关机", "on": "开机"}
		# 首先判断当前时间是否在配置文件设置的时间范围内
		if self.get_config_final_time(hour, delayed):
			# 如果在时间范围内，则判断当前状态是否符合配置文件设置的操作
			if self.debug:
				logger.debug(f"正在{op_dict_name[action]}: {Id}")
			if os.system(f"lxc-{op_dict[action]} {Id}") == 0:
				logger.info(f"{op_dict_name[action]}成功: {Id}")
				self.result_log(f"{datetime.now()} - {Id} - {op_dict_name[action]}成功\n")
			else:
				logger.error(f"{op_dict_name[action]}失败: {Id}")
				self.result_log(f"{datetime.now()} - {Id} - {op_dict_name[action]}失败\n")

	def list_config(self):
		"""
		读取列表主机配置文件
		:return:
		"""
		# 判断文件夹是否存在且是否有文件
		if path.isdir(self.config_list):
			fl = list_files(self.config_list)
			if fl:
				if self.debug:
					# 打印文件数量
					logger.debug(f"当前配置文件数量: {len(fl)}")
				for file in fl:
					self.read_list(file=file)
					# 获取需要操作的主机ID列表
					data_list = self.txt_config_list["list"]
					host_list = data_list["list"]
					hour_ = data_list["hour"]
					delayed_ = data_list["delayed"]
					action_ = data_list["action"]
					for i in host_list:
						if self.get_config_final_time(hour_, delayed_):
							self.pve_op_ct_list(Id=i, hour=hour_, delayed=delayed_, action=action_)
						else:
							if self.debug:
								logger.debug(f"当前时间不在配置范围： {self.current_hour}:{self.current_min}")
			else:
				logger.warning(f"配置文件目录为空: {self.config_list}")
		else:
			if self.debug:
				logger.debug(f"配置文件目录不存在: {self.config_list}")

	def start(self):
		while True:
			self.restart()
			self.list_config()


def install():
	"""
	安装
	:return:
	"""
	logger.info("正在进行安装/更新...")
	url = 'https://gitee.com/liumou_site/pvet/raw/master/model/lxc/list.toml'
	url_restart = 'https://gitee.com/liumou_site/pvet/raw/master/model/lxc/restart.toml'
	os.system("mkdir -p /opt/apps/com.liumou.pve.lxc/")
	restart = "/opt/apps/com.liumou.pve.lxc/demo_restart.toml"
	demo = "/opt/apps/com.liumou.pve.lxc/demo_list.toml"
	url_dict = {demo: url, restart: url_restart}
	for file in url_dict:
		if path.exists(file):
			logger.info("模板配置文件已存在")
			os.system(f"rm -f {file}")
		res = subprocess.getstatusoutput(f'wget -O {file} {url_dict[file]}')
		if res[0] == 0:
			logger.info(f"[ {file} ]配置文件下载成功")
		else:
			logger.info(f"[ {file} ]配置文件下载失败")
			print(res[1])
			exit(1)
	lxc_url = 'https://gitee.com/liumou_site/pvet/raw/master/model/lxc/lxc.py'
	lxc_bin = path.join("/usr/bin/com.liumou.pve.lxc")
	os.system(f"rm -f {lxc_bin}")
	res = subprocess.getstatusoutput(f'wget -O {lxc_bin} {lxc_url}')
	if res[0] == 0:
		os.system(f"chmod 0755 {lxc_bin}")
		logger.info("安装/更新成功")
		service()
		exit(0)
	else:
		logger.error("安装/更新失败")
		print(res[1])
		exit(2)


def service():
	"""

	:return:
	"""
	if not path.isfile("/usr/bin/com.liumou.pve.lxc"):
		logger.error(f"执行文件不存在: /usr/bin/com.liumou.pve.lxc")
		exit(1)
	txt = f"""[Unit]
Description=This is a PVE timing tool
After=network.target

[Service]
Type=simple
User=root
ExecStart=/usr/bin/com.liumou.pve.lxc
StandardOutput=journal
KillMode=control-group
Restart=on-failure
RestartPreventExitStatus=255
StandardError=inherit
SyslogIdentifier=com.liumou.pve.lxc

[Install]
WantedBy=multi-user.target
Alias=com.liumou.pve.lxc.service
"""
	try:
		with open(file="/etc/systemd/system/com.liumou.pve.lxc.service", mode='w', encoding='utf-8') as w:
			w.write(txt)
	except Exception as e:
		logger.error(f"写入服务文件失败: {e}")
		exit(1)
	if os.system("systemctl daemon-reload") == 0:
		os.system("systemctl enable com.liumou.pve.lxc.service")
		os.system("systemctl restart com.liumou.pve.lxc.service")
		logger.info("服务启动成功")
		exit(0)
	else:
		logger.error("服务启动失败")
		exit(2)


if __name__ == '__main__':
	conf_default = "/opt/apps/com.liumou.pve.lxc/list.toml"
	conf_default_restart_dir = "/opt/apps/com.liumou.pve.lxc/config_restart/"
	conf_default_dir = "/opt/apps/com.liumou.pve.lxc/config_restart/"
	os.system(f"mkdir -p /opt/apps/com.liumou.pve.lxc/log/ {conf_default_dir}")

	arg = ArgumentParser(description='当前脚本版本: 2.1', prog="PVE容器定时开关机")
	h = f"设置单机配置文件路径,默认: {conf_default}"
	arg.add_argument('-c', '--conf', type=str, help=h, default=conf_default, required=False)
	h = f"设置更多配置文件所在目录(此目录下所有.toml结尾的文件都会被读取),默认: {conf_default_dir}"
	arg.add_argument('-cd', '--conf_dir', type=str, help=h, default=conf_default_dir, required=False)
	arg.add_argument('-install', action='store_true', help='安装/更新服务')
	arg.add_argument('-f', action='store_true', help='强制更新配置文件')
	arg.add_argument('-debug', action='store_true', help='开启调试模式')
	args = arg.parse_args()
	args_conf = args.conf
	args_install = args.install
	args_debug = args.debug
	args_f = args.f
	args_conf_restart = args.conf_dir
	if args_debug:
		logger.info("开启了调试模式")
	if args_install:
		logger.info("开启了安装/更新功能")
		install()
	if not path.exists(args_conf):
		logger.error(f"配置文件不存在: {args_conf}")
		exit(1)
	c = Lxc(config_restart_dir=args_conf, debug=args_debug, config_list=args_conf)
	c.start()
